from langchain_community.llms import Tongyi
from langchain.prompts import FewShotChatMessagePromptTemplate, ChatPromptTemplate
import os

os.environ["DASHSCOPE_API_KEY"] = 'sk-c44402d7a12c41299bb716af8d7e8bac'

human_prompt_content = """
    格式化以下信息：
        姓名---{user_name}
        性别---{user_sex}
        年龄---{user_age}
        工作---{user_job}
        城市---{user_city}
"""

ai_prompt_content = """
    姓名:{f_user_name}
    性别:{f_user_sex}
    年龄:{f_user_age}
    工作:{f_user_job}
    城市:{f_user_city}
"""

example_prompt = ChatPromptTemplate.from_messages(
    [
        ('system', human_prompt_content),
        ('ai', ai_prompt_content)
    ]
)

examples = [
    {
        'user_name': '张三',
        'user_sex': '男',
        'user_age': '25',
        'user_job': '后端',
        'user_city': '武汉',
        'f_user_name': '张三先生',
        'f_user_sex': '男性',
        'f_user_age': '25岁',
        'f_user_job': '后端开发工程师',
        'f_user_city': '湖北省武汉市'
    },
    {
        'user_name': '李四',
        'user_sex': '女',
        'user_age': '20',
        'user_job': '前端',
        'user_city': '广州',
        'f_user_name': '李四女士',
        'f_user_sex': '女性',
        'f_user_age': '20岁',
        'f_user_job': '前端开发工程师',
        'f_user_city': '广东省广州市'
    }
]

few_shot_prompt = FewShotChatMessagePromptTemplate(
    example_prompt=example_prompt,
    examples=examples
)

human_prompt_content01 = """
    格式化以下信息：
        姓名---{user_name}
        性别---{user_sex}
        年龄---{user_age}
        工作---{user_job}
        城市---{user_city}
"""

prompt_template = ChatPromptTemplate.from_messages(
    [
        few_shot_prompt,
        ('human', human_prompt_content)
    ]
)

prompt = prompt_template.invoke({'user_name': '王五', 'user_sex': '男', 'user_age': '30', 'user_job': 'ai', 'user_city': '武汉'})

llm = Tongyi()

response = llm.invoke(prompt)
print(response)
